Skip to content

Implement aten.stft#2645

Merged
titaiwangms merged 3 commits intomicrosoft:mainfrom
moatom:add_stft
Nov 9, 2025
Merged

Implement aten.stft#2645
titaiwangms merged 3 commits intomicrosoft:mainfrom
moatom:add_stft

Conversation

@moatom
Copy link
Contributor

@moatom moatom commented Oct 19, 2025 •
edited
Loading

Fixed pytorch/pytorch#147052

$ python -m pytest tests/function_libs/torch_lib/ops_test.py -k ops_aten_stft
====================================================================================================================================================================================================== test session starts ======================================================================================================================================================================================================
platform linux -- Python 3.13.1, pytest-8.4.1, pluggy-1.6.0
Using --randomly-seed=371864411
rootdir: /home/moatom/github/onnxscript
configfile: pyproject.toml
plugins: randomly-3.16.0, xdist-3.8.0, subtests-0.14.2, cov-6.2.1, hypothesis-6.138.2
collected 2158 items / 2154 deselected / 4 selected                                                                                                                                                                                                                                                                                                                                                                             

tests/function_libs/torch_lib/ops_test.py s..x                                                                                                                                                                                                                                                                                                                                                                     [100%]

======================================================================================================================================================================================================= warnings summary ========================================================================================================================================================================================================
onnxscript/converter.py:457: 429 warnings
tests/function_libs/torch_lib/ops_test.py: 15 warnings
  /home/moatom/github/onnxscript/onnxscript/converter.py:457: DeprecationWarning: Expression.__init__ got an unexpected keyword argument 'lineno'. Support for arbitrary keyword arguments is deprecated and will be removed in Python 3.15.
    expr = ast.Expression(expr, lineno=expr.lineno, col_offset=expr.col_offset)

onnxscript/converter.py:457: 429 warnings
tests/function_libs/torch_lib/ops_test.py: 15 warnings
  /home/moatom/github/onnxscript/onnxscript/converter.py:457: DeprecationWarning: Expression.__init__ got an unexpected keyword argument 'col_offset'. Support for arbitrary keyword arguments is deprecated and will be removed in Python 3.15.
    expr = ast.Expression(expr, lineno=expr.lineno, col_offset=expr.col_offset)

tests/function_libs/torch_lib/ops_test.py::TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__ops_aten_stft_cpu_float32
tests/function_libs/torch_lib/ops_test.py::TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__ops_aten_stft_cpu_float32
tests/function_libs/torch_lib/ops_test.py::TestOutputConsistencyFullGraphCPU::test_output_match_opinfo__ops_aten_stft_cpu_float32
  /home/moatom/github/onnxscript/tests/function_libs/torch_lib/ops_test_common.py:329: DeprecationWarning: __array__ implementation doesn't accept a copy keyword, so passing copy=False failed. __array__ must implement 'dtype' and 'copy' keyword arguments. To learn more, see the migration guide https://numpy.org/devdocs/numpy_2_0_migration_guide.html#adapting-to-changes-in-the-copy-keyword
    value = np.array(value.cpu())

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
==================================================================================================================================================================================================== short test summary info ====================================================================================================================================================================================================
SKIPPED [1] tests/function_libs/torch_lib/ops_test.py:101: Traced functions does not have a function proto
=================================================================================================================================================================== 2 passed, 1 skipped, 2154 deselected, 1 xfailed, 891 warnings, 7 subtests passed in 4.42s ===================================================================================================================================================================

@moatom moatom changed the title Add stft Implement aten.stft Oct 19, 2025
@moatom
Copy link
Contributor Author

moatom commented Oct 19, 2025

@microsoft-github-policy-service agree

@moatom
Copy link
Contributor Author

moatom commented Oct 19, 2025

pytorch/pytorch#147052 (comment)

It would be nice to decay to onnx.STFT node in the graph if options of a given call allow this...

I think this PR’s implementation already uses onnx.STFT as its core: https://github.com/microsoft/onnxscript/pull/2645/files#diff-773a426eaaf0eb2a1ed6039f2cd50613d2db7d61c239453545413499b2a15776R8194

result = op.STFT(signal, frame_step_const, window, frame_length_const, onesided=onesided)

So I don’t think an decay to onnx.STFT is necessary in this case. 🤔

@moatom
Copy link
Contributor Author

moatom commented Oct 19, 2025 •
edited
Loading

@justinchuby

Hi!

pytorch/pytorch#147052 (comment)

I think it can be simplified quite a bit.

Could you elaborate a bit on this comment at your convenience? (If you think this simplification is necessary.)

@moatom moatom marked this pull request as ready for review October 19, 2025 15:00
@justinchuby justinchuby requested review from Copilot and titaiwangms and removed request for titaiwangms October 20, 2025 16:17
Copilot AI reviewed Oct 20, 2025
View reviewed changes
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull Request Overview

This PR implements the aten::stft (Short-Time Fourier Transform) operator to resolve issue #147052. The implementation includes handling for various optional parameters like hop_length, win_length, window, normalized, onesided, and return_complex.

Key changes:

  • Added STFT operator implementation with helper functions for batch dimension handling, window centering, and FFT normalization
  • Registered the operator in test data with appropriate tolerance settings and xfail for float16 dtype

Reviewed Changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File Description
onnxscript/function_libs/torch_lib/ops/core.py Implements aten_stft and five helper functions for STFT processing
tests/function_libs/torch_lib/ops_test_data.py Registers the new operator in test suite with tolerance and xfail configuration

@titaiwangms titaiwangms self-assigned this Oct 24, 2025
@codecov
Copy link

codecov bot commented Oct 24, 2025 •
edited
Loading

Codecov Report

❌ Patch coverage is 65.45455% with 19 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.11%. Comparing base (a1be5c8) to head (29ba6b9).
⚠️ Report is 1 commits behind head on main.

Files with missing lines Patch % Lines
onnxscript/function_libs/torch_lib/ops/core.py 65.45% 15 Missing and 4 partials ⚠️
Additional details and impacted files 
@@           Coverage Diff           @@
##             main    #2645   +/-   ##
=======================================
  Coverage   70.11%   70.11%           
=======================================
  Files         225      225           
  Lines       27083    27138   +55     
  Branches     2721     2730    +9     
=======================================
+ Hits        18990    19029   +39     
- Misses       7153     7165   +12     
- Partials      940      944    +4

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

titaiwangms
titaiwangms reviewed Oct 24, 2025
View reviewed changes
titaiwangms
titaiwangms reviewed Oct 28, 2025
View reviewed changes
Copy link
Contributor

@titaiwangms titaiwangms left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Some minor comments and this should be the final pass.
cc @justinchuby I remember you said something stft implementation was wrong. This is heavily following the original one. Would you care share about your comments?

@titaiwangms titaiwangms added the module: torchlib Related to the torch/aten function lib in development label Oct 28, 2025
@justinchuby
Copy link
Collaborator

justinchuby commented Oct 28, 2025

Thanks - I need to read this more closely to remind myself what was going on then

justinchuby
justinchuby approved these changes Oct 29, 2025
View reviewed changes
Copy link
Collaborator

@justinchuby justinchuby left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks. I don't have a better idea

@titaiwangms
Copy link
Contributor

titaiwangms commented Oct 29, 2025

@moatom Could you do a final pass, and make sure lint is passing.

@moatom
Copy link
Contributor Author

moatom commented Nov 1, 2025

@titaiwangms I resolved all of them. Could you please confirm when you have a moment?

justinchuby
justinchuby reviewed Nov 1, 2025
View reviewed changes
@justinchuby
Copy link
Collaborator

justinchuby commented Nov 1, 2025

If possible, could you help add a few tests to https://github.com/microsoft/onnxscript/blob/main/tests/function_libs/torch_lib/e2e_ops_tests.py ? Thanks a lot!

@moatom moatom force-pushed the add_stft branch 2 times, most recently from 76ffeca to c0fd969 Compare November 2, 2025 14:14
@moatom
Copy link
Contributor Author

moatom commented Nov 2, 2025 •
edited
Loading

@justinchuby

If possible, could you help add a few tests to https://github.com/microsoft/onnxscript/blob/main/tests/function_libs/torch_lib/e2e_ops_tests.py ?

DONE

@github-project-automation github-project-automation bot moved this from Todo to Done in ONNX Script Review Board Nov 9, 2025
@titaiwangms titaiwangms enabled auto-merge (squash) November 9, 2025 05:49
@moatom
Copy link
Contributor Author

moatom commented Nov 9, 2025 •
edited
Loading

Lint / Enforce style (pull_request)Failing after 5m

I am sorry. I will fix it now.

auto-merge was automatically disabled November 9, 2025 06:02

Head branch was pushed to by a user without write access

@moatom
Copy link
Contributor Author

moatom commented Nov 9, 2025

@titaiwangms

Lint / Enforce style (pull_request)Failing after 5m

I am sorry. I will fix it now.

DONE.

@titaiwangms titaiwangms enabled auto-merge (squash) November 9, 2025 17:34
@titaiwangms titaiwangms merged commit 10e541e into microsoft:main Nov 9, 2025
32 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

Copilot code review Copilot Copilot left review comments

@titaiwangms titaiwangms titaiwangms approved these changes

@justinchuby justinchuby Awaiting requested review from justinchuby

Assignees

@justinchuby justinchuby

Labels

module: torchlib Related to the torch/aten function lib in development

Projects

ONNX Script Review Board
Status: Done

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[ONNX] Implement aten.stft

3 participants

@moatom @justinchuby @titaiwangms